import matplotlib
from matplotlib import pyplot
import minuit
import numpy

import csbbelle.fit.tex

class TwoGaussPlusPolynomialFit(GaussPlusPolynomialFit):
    tex_labels = {'A1' : r'A_1',
                  'mu1': r'\mu_1',
                  'sigma1': r'\sigma_1',
                  'A2' : r'A_2',
                  'mu2': r'\mu_2',
                  'sigma2': r'\sigma_2',
                  'p0': r'p_0',
                  'p1': r'p_1',
                  'p2': r'p_2',
                  'p3': r'p_3',
                  'p4': r'p_4',
                  'p5': r'p_5',
                  'p6': r'p_6',
                  'p7': r'p_7',
                  'p8': r'p_8',
                  'p9': r'p_9',
                }
     
    tex_header=r'$A_1\exp\left[\frac{(x - \mu_1)^2}{\sigma_1^2}\right] ' + \
        r'+ A_2\exp\left[\frac{(x - \mu_2)^2}{\sigma_2^2}\right]'

    signal_fields = "A1 mu1 sigma1 A2 mu2 sigma2".split()

    def _signal(self, x, A1, mu1, sigma1, A2, mu2, sigma2):
        return A1*numpy.exp(-0.5*((x - mu1)/sigma1)**2) + \
            A2*numpy.exp(-0.5*((x - mu2)/sigma2)**2)

    def _bg(self, x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        return numpy.polynomial.chebyshev.chebval(x, [p0, p1, p2, p3, p4, p5, p6, p7, p8, p9])

    def _model(self, x, A1, mu1, sigma1, A2, mu2, sigma2,
               p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        return self._signal(x, A1, mu1, sigma1, A2, mu2, sigma2) + \
            self._bg(x, p0, p1, p2, p3, p4, p5, p6, p7, p8, p9)
    
    def _chi2(self, A1, mu1, sigma1, A2, mu2, sigma2, 
              p0, p1, p2, p3, p4, p5, p6, p7, p8, p9):
        chi = (self._model(self.fit_x, A1, mu1, sigma1, A2, mu2, sigma2, 
                           p0, p1, p2, p3, p4, p5, p6, p7, p8, p9) 
               - self.fit_y) / self.fit_yunc

        chisq = numpy.sum(chi**2)

        if self.loudness > 0 and self.ncalls % 100 == 0:
            print self.ncalls, chisq
            
        return chisq

    def signal(self):
        return self._signal(self.fit_x,
                            self.m.values['A1'], 
                            self.m.values['mu1'],
                            self.m.values['sigma1'],
                            self.m.values['A2'], 
                            self.m.values['mu2'],
                            self.m.values['sigma2'])

    def bg(self):
        return self._bg(self.fit_x,
                        self.m.values['p0'],
                        self.m.values['p1'],
                        self.m.values['p2'],
                        self.m.values['p3'],
                        self.m.values['p4'],
                        self.m.values['p5'],
                        self.m.values['p6'],
                        self.m.values['p7'],
                        self.m.values['p8'],
                        self.m.values['p9'],
                        )


    def release_all(self):
        self.release_signal()
        self.release_background()
        

            
